package org.jeecg.modules.airag.llm.tools;

import org.jeecg.modules.airag.flow.component.enhance.IAiRagEnhanceJava;
import org.neo4j.driver.*;
import org.neo4j.driver.Record;
import org.springframework.stereotype.Component;

import java.util.*;
import java.util.stream.Collectors;

import static org.neo4j.driver.Values.parameters;

@Component("Neo4jQuery")
public class Neo4jQuery implements IAiRagEnhanceJava {

    // Neo4j 配置
    private static final String NEO4J_URI = "bolt://192.168.31.83:7687";
    private static final String NEO4J_USER = "neo4j";
    private static final String NEO4J_PASSWORD = "cxy052166";

    // 创建Neo4j驱动
    private Driver driver =  GraphDatabase.driver(NEO4J_URI,
            AuthTokens.basic(NEO4J_USER, NEO4J_PASSWORD));

    @Override
    public Map<String, Object> process(Map<String, Object> map) {
        String queryCypher = (String) map.get("queryCypher");
        String queryCypher1 = (String) map.get("queryCypher1");
        System.out.println("queryCypher1 = " + queryCypher1);
        if (queryCypher1 != null){
            queryCypher = queryCypher1;
        }
//        // 清理查询字符串
        queryCypher = cleanQueryCypher(queryCypher);
//        Cypher
        System.out.println("queryCypher = " + queryCypher);
//        queryCypher.replace("Cypher:", "");
        if (queryCypher == null || queryCypher.trim().isEmpty()) {
            return Map.of("error", "queryCypher is required");
        }
        try {
            // 执行查询
            List<Map<String, Object>> results = executeQuery(queryCypher);

            if (results.isEmpty()) {
                return Map.of("error", "未查询到有效数据");
            }

            // 生成Markdown表格
//            String markdownTable = generateMarkdownTable(results);
            return Map.of("result", results);

        } catch (Exception e) {
            return Map.of("error", e.getMessage());
        }
    }
    /**
     * 清理查询字符串，去除markdown标记和多余字符
     */
    private String cleanQueryCypher(String queryCypher) {
        if (queryCypher == null) {
            return null;
        }

        // 去除markdown标记和多余字符
        String cleaned = queryCypher.replace("```cypher\n", "")
                .replace("```text\n", "")
                .replace("\n<<EOF>>", "")
                .replace("```", " ")
                .replace("：", ":")
                .trim()
                .replace("\n", " ")
                .replaceAll("\\s+", " "); // 多个空格替换为一个

        return cleaned;
    }

    private Object simplifyValue(Value value) {
        if (value == null || value.isNull()) return null;

        // 基本类型
        try { return value.asBoolean(); } catch (Exception e) {}
        try { return value.asLong(); } catch (Exception e) {}
        try { return value.asDouble(); } catch (Exception e) {}
        try { return value.asString(); } catch (Exception e) {}

        // 列表
        try {
            return value.asList(this::simplifyValue);
        } catch (Exception e) {}

        // 映射（属性、对象）
        try {
            return value.asMap().entrySet().stream()
                    .collect(Collectors.toMap(
                            Map.Entry::getKey,
                            entry -> entry.getValue().toString()
                    ));
        } catch (Exception e) {}

        // 兜底：转字符串（Node、Relationship、Path 等复杂类型）
        return value.toString();
    }

    /**
     * 执行Cypher查询并返回结果
     */
    private List<Map<String, Object>> executeQuery(String query) {
        try (Session session = driver.session()) {
            Result result = session.run(query, parameters());

            List<Map<String, Object>> rows = new ArrayList<>();
            while (result.hasNext()) {
                Record record = result.next();
                Map<String, Object> row = new LinkedHashMap<>();
                for (String key : record.keys()) {
                    row.put(key, simplifyValue(record.get(key)));
                }
                rows.add(row);
            }
            return rows;
        } catch (Exception e) {
            throw new RuntimeException("执行 query 失败: " + query, e);
        }
    }

    /**
     * 生成Markdown表格
     */
    private String generateMarkdownTable(List<Map<String, Object>> results) {
        if (results.isEmpty()) {
            return "";
        }

        // 获取列名
        List<String> columns = results.get(0).keySet().stream().collect(Collectors.toList());

        // 构建表头
        StringBuilder table = new StringBuilder();
        table.append("| ").append(String.join(" | ", columns)).append(" |\n");

        // 构建分隔线
        table.append("| ").append(columns.stream()
                .map(col -> "---")
                .collect(Collectors.joining(" | "))).append(" |\n");

        // 构建行数据
        for (Map<String, Object> row : results) {
            List<String> values = columns.stream()
                    .map(col -> String.valueOf(row.get(col)))
                    .collect(Collectors.toList());
            table.append("| ").append(String.join(" | ", values)).append(" |\n");
        }

        return table.toString();
    }

}
